/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.rel;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;

import static org.apache.calcite.avatica.util.DateTimeUtils.MILLIS_PER_DAY;



import java.io.IOException;

import java.util.Iterator;

import java.util.List;

import java.util.Map;

import java.util.Queue;

import java.util.concurrent.ConcurrentHashMap;

import java.util.concurrent.ConcurrentLinkedQueue;

import java.util.stream.Collectors;

import javax.annotation.Nullable;

import com.bff.gaia.unified.runners.gaia.GaiaPipelineOptions;



import com.bff.gaia.unified.sdk.Pipeline;

import com.bff.gaia.unified.sdk.PipelineResult;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.CalciteUtils.CharType;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.CalciteUtils.DateType;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.CalciteUtils.TimeType;

import com.bff.gaia.unified.sdk.metrics.Counter;

import com.bff.gaia.unified.sdk.metrics.MetricNameFilter;

import com.bff.gaia.unified.sdk.metrics.MetricQueryResults;

import com.bff.gaia.unified.sdk.metrics.MetricResult;

import com.bff.gaia.unified.sdk.metrics.Metrics;

import com.bff.gaia.unified.sdk.metrics.MetricsFilter;

import com.bff.gaia.unified.sdk.options.ApplicationNameOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptionsFactory;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.ParDo;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.PValue;

import com.bff.gaia.unified.sdk.values.Row;

import org.apache.calcite.adapter.enumerable.EnumerableRel;

import org.apache.calcite.adapter.enumerable.EnumerableRelImplementor;

import org.apache.calcite.adapter.enumerable.PhysType;

import org.apache.calcite.adapter.enumerable.PhysTypeImpl;

import org.apache.calcite.linq4j.Enumerable;

import org.apache.calcite.linq4j.Linq4j;

import org.apache.calcite.linq4j.tree.BlockBuilder;

import org.apache.calcite.linq4j.tree.Expression;

import org.apache.calcite.linq4j.tree.Expressions;

import org.apache.calcite.plan.ConventionTraitDef;

import org.apache.calcite.plan.RelOptCluster;

import org.apache.calcite.plan.RelOptCost;

import org.apache.calcite.plan.RelOptPlanner;

import org.apache.calcite.plan.RelTraitSet;

import org.apache.calcite.rel.RelNode;

import org.apache.calcite.rel.convert.ConverterImpl;

import org.apache.calcite.rel.metadata.RelMetadataQuery;

import org.apache.calcite.rel.type.RelDataType;

import com.bff.gaia.unified.sdk.runners.TransformHierarchy;

import org.joda.time.Duration;

import org.joda.time.ReadableInstant;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



/** UnifiedRelNode to replace a {@code Enumerable} node. */

public class UnifiedEnumerableConverter extends ConverterImpl implements EnumerableRel {

  private static final Logger LOG = LoggerFactory.getLogger(UnifiedEnumerableConverter.class);



  public UnifiedEnumerableConverter(RelOptCluster cluster, RelTraitSet traits, RelNode input) {

    super(cluster, ConventionTraitDef.INSTANCE, traits, input);

  }



  @Override

  public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {

    return new UnifiedEnumerableConverter(getCluster(), traitSet, sole(inputs));

  }



  @Override

  public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {

    // This should always be a last resort.

    return planner.getCostFactory().makeHugeCost();

  }



  @Override

  public Result implement(EnumerableRelImplementor implementor, Prefer prefer) {

    final BlockBuilder list = new BlockBuilder();

    final RelDataType rowType = getRowType();

    final PhysType physType =

        PhysTypeImpl.of(implementor.getTypeFactory(), rowType, prefer.preferArray());

    final Expression node = implementor.stash((UnifiedRelNode) getInput(), UnifiedRelNode.class);

    list.add(Expressions.call(UnifiedEnumerableConverter.class, "toEnumerable", node));

    return implementor.result(physType, list.toBlock());

  }



  public static Enumerable<Object> toEnumerable(UnifiedRelNode node) {

    final ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();

    try {

      Thread.currentThread().setContextClassLoader(UnifiedEnumerableConverter.class.getClassLoader());

      final PipelineOptions options = createPipelineOptions(node.getPipelineOptions());

      return toEnumerable(options, node);

    } finally {

      Thread.currentThread().setContextClassLoader(originalClassLoader);

    }

  }



  public static List<Row> toRowList(UnifiedRelNode node) {

    final ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();

    try {

      Thread.currentThread().setContextClassLoader(UnifiedEnumerableConverter.class.getClassLoader());

      final PipelineOptions options = createPipelineOptions(node.getPipelineOptions());

      return toRowList(options, node);

    } finally {

      Thread.currentThread().setContextClassLoader(originalClassLoader);

    }

  }



  public static PipelineOptions createPipelineOptions(Map<String, String> map) {

    final String[] args = new String[map.size()];

    int i = 0;

    for (Map.Entry<String, String> entry : map.entrySet()) {

      args[i++] = "--" + entry.getKey() + "=" + entry.getValue();

    }

    PipelineOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().create();

    options.as(ApplicationNameOptions.class).setAppName("UnifiedSql");

    return options;

  }



  static List<Row> toRowList(PipelineOptions options, UnifiedRelNode node) {

    if (node instanceof UnifiedIOSinkRel) {

      throw new UnsupportedOperationException("Does not support UnifiedIOSinkRel in toRowList.");

    } else if (isLimitQuery(node)) {

      throw new UnsupportedOperationException("Does not support queries with LIMIT in toRowList.");

    }

    return collectRows(options, node).stream().collect(Collectors.toList());

  }



  static Enumerable<Object> toEnumerable(PipelineOptions options, UnifiedRelNode node) {

    if (node instanceof UnifiedIOSinkRel) {

      return count(options, node);

    } else if (isLimitQuery(node)) {

      return limitCollect(options, node);

    }

    return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues((collectRows(options, node))));

  }



  private static PipelineResult limitRun(

      PipelineOptions options,

      UnifiedRelNode node,

      DoFn<Row, Void> doFn,

      Queue<Row> values,

      int limitCount) {

    options.as(GaiaPipelineOptions.class);



    Pipeline pipeline = Pipeline.create(options);

    PCollection<Row> resultCollection = UnifiedSqlRelUtils.toPCollection(pipeline, node);

    resultCollection.apply(ParDo.of(doFn));



    PipelineResult result = pipeline.run();



    PipelineResult.State state;

    while (true) {

      // Check pipeline state in every second

      state = result.waitUntilFinish(Duration.standardSeconds(1));

      if (state != null && state.isTerminal()) {

        break;

      }



      try {

        if (values.size() >= limitCount) {

          result.cancel();

          break;

        }

      } catch (IOException e) {

        LOG.warn(e.toString());

        break;

      }

    }



    return result;

  }



  private static void runCollector(PipelineOptions options, UnifiedRelNode node) {

    Pipeline pipeline = Pipeline.create(options);

    PCollection<Row> resultCollection = UnifiedSqlRelUtils.toPCollection(pipeline, node);

    resultCollection.apply(ParDo.of(new Collector()));

    PipelineResult result = pipeline.run();

    result.waitUntilFinish();

  }



  private static Queue<Row> collectRows(PipelineOptions options, UnifiedRelNode node) {

    long id = options.getOptionsId();

    Queue<Row> values = new ConcurrentLinkedQueue<>();



    checkArgument(

        options

            .getRunner()

            .getCanonicalName()

            .equals("com.bff.gaia.unified.runners.direct.DirectRunner"),

        "collectRowList is only available in direct runner.");



    Collector.globalValues.put(id, values);



    runCollector(options, node);



    Collector.globalValues.remove(id);

    return values;

  }



  private static Enumerable<Object> limitCollect(PipelineOptions options, UnifiedRelNode node) {

    long id = options.getOptionsId();

    ConcurrentLinkedQueue<Row> values = new ConcurrentLinkedQueue<>();



    checkArgument(

        options

            .getRunner()

            .getCanonicalName()

            .equals("com.bff.gaia.unified.runners.direct.DirectRunner"),

        "SELECT without INSERT is only supported in DirectRunner in SQL Shell.");



    int limitCount = getLimitCount(node);



    Collector.globalValues.put(id, values);

    limitRun(options, node, new Collector(), values, limitCount);

    Collector.globalValues.remove(id);



    // remove extra retrieved values

    while (values.size() > limitCount) {

      values.remove();

    }



    return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues(values));

  }



  private static class Collector extends DoFn<Row, Void> {



    // This will only work on the direct runner.

    private static final Map<Long, Queue<Row>> globalValues = new ConcurrentHashMap<>();



    @Nullable

	private volatile Queue<Row> values;



    @StartBundle

    public void startBundle(DoFn.StartBundleContext context) {

      long id = context.getPipelineOptions().getOptionsId();

      values = globalValues.get(id);

    }



    @ProcessElement

    public void processElement(ProcessContext context) {

      values.add(context.element());

    }

  }



  private static List<Object> rowToAvaticaAndUnboxValues(Queue<Row> values) {

    return values.stream()

        .map(

            row -> {

              Object[] objects = rowToAvatica(row);

              if (objects.length == 1) {

                // if objects.length == 1, that means input Row contains only 1 column/element,

                // then an Object instead of Object[] should be returned because of

                // CalciteResultSet's behaviour that tries to convert one column row to an Object.

                return objects[0];

              } else {

                return objects;

              }

            })

        .collect(Collectors.toList());

  }



  private static Object[] rowToAvatica(Row row) {

    Schema schema = row.getSchema();

    Object[] convertedColumns = new Object[schema.getFields().size()];

    int i = 0;

    for (Schema.Field field : schema.getFields()) {

      convertedColumns[i] = fieldToAvatica(field.getType(), row.getValue(i));

      ++i;

    }

    return convertedColumns;

  }



  private static Object fieldToAvatica(Schema.FieldType type, Object unifiedValue) {

    if (unifiedValue == null) {

      return null;

    }



    switch (type.getTypeName()) {

      case LOGICAL_TYPE:

        String logicalId = type.getLogicalType().getIdentifier();

        if (logicalId.equals(TimeType.IDENTIFIER)) {

          return (int) ((ReadableInstant) unifiedValue).getMillis();

        } else if (logicalId.equals(DateType.IDENTIFIER)) {

          return (int) (((ReadableInstant) unifiedValue).getMillis() / MILLIS_PER_DAY);

        } else if (logicalId.equals(CharType.IDENTIFIER)) {

          return unifiedValue;

        } else {

          throw new IllegalArgumentException("Unknown DateTime type " + logicalId);

        }

      case DATETIME:

        return ((ReadableInstant) unifiedValue).getMillis();

      case BYTE:

      case INT16:

      case INT32:

      case INT64:

      case DECIMAL:

      case FLOAT:

      case DOUBLE:

      case STRING:

      case BOOLEAN:

      case BYTES:

        return unifiedValue;

      case ARRAY:

        return ((List<?>) unifiedValue)

            .stream()

                .map(elem -> fieldToAvatica(type.getCollectionElementType(), elem))

                .collect(Collectors.toList());

      case MAP:

        return ((Map<?, ?>) unifiedValue)

            .entrySet().stream()

                .collect(

                    Collectors.toMap(

                        entry -> entry.getKey(),

                        entry ->

                            fieldToAvatica(type.getCollectionElementType(), entry.getValue())));

      case ROW:

        // TODO: needs to be a Struct

        return unifiedValue;

      default:

        throw new IllegalStateException(

            String.format("Unreachable case for Unified typename %s", type.getTypeName()));

    }

  }



  private static Enumerable<Object> count(PipelineOptions options, UnifiedRelNode node) {

    Pipeline pipeline = Pipeline.create(options);

    UnifiedSqlRelUtils.toPCollection(pipeline, node).apply(ParDo.of(new RowCounter()));

    PipelineResult result = pipeline.run();



    long count = 0;

    if (!containsUnboundedPCollection(pipeline)) {

      result.waitUntilFinish();

      MetricQueryResults metrics =

          result

              .metrics()

              .queryMetrics(

                  MetricsFilter.builder()

                      .addNameFilter(MetricNameFilter.named(UnifiedEnumerableConverter.class, "rows"))

                      .build());

      Iterator<MetricResult<Long>> iterator = metrics.getCounters().iterator();

      if (iterator.hasNext()) {

        count = iterator.next().getAttempted();

      }

    }

    return Linq4j.singletonEnumerable(count);

  }



  private static class RowCounter extends DoFn<Row, Void> {

    final Counter rows = Metrics.counter(UnifiedEnumerableConverter.class, "rows");



    @ProcessElement

    public void processElement(DoFn.ProcessContext context) {

      rows.inc();

    }

  }



  private static boolean isLimitQuery(UnifiedRelNode node) {

    return (node instanceof UnifiedSortRel && ((UnifiedSortRel) node).isLimitOnly())

        || (node instanceof UnifiedCalcRel && ((UnifiedCalcRel) node).isInputSortRelAndLimitOnly());

  }



  private static int getLimitCount(UnifiedRelNode node) {

    if (node instanceof UnifiedSortRel) {

      return ((UnifiedSortRel) node).getCount();

    } else if (node instanceof UnifiedCalcRel) {

      return ((UnifiedCalcRel) node).getLimitCountOfSortRel();

    }



    throw new RuntimeException(

        "Cannot get limit count from RelNode tree with root " + node.getRelTypeName());

  }



  private static boolean containsUnboundedPCollection(Pipeline p) {

    class BoundednessVisitor extends Pipeline.PipelineVisitor.Defaults {

      PCollection.IsBounded boundedness = PCollection.IsBounded.BOUNDED;



      @Override

      public void visitValue(PValue value, TransformHierarchy.Node producer) {

        if (value instanceof PCollection) {

          boundedness = boundedness.and(((PCollection) value).isBounded());

        }

      }

    }



    BoundednessVisitor visitor = new BoundednessVisitor();

    p.traverseTopologically(visitor);

    return visitor.boundedness == PCollection.IsBounded.UNBOUNDED;

  }

}